{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 03: weak imposition of Dirichlet BCs by a Lagrange multiplier (interface problem)\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{cases}\n",
    "-\\Delta u = f, & \\text{in } \\Omega,\\\\\n",
    " u   = g, & \\text{on } \\partial\\Omega,\n",
    "\\end{cases}$$\n",
    "\n",
    "where $\\Omega$ is the unit ball in 2D, using a domain decomposition approach for $\\Omega = \\Omega_1 \\cup \\Omega_2$, and introducing a lagrange multiplier to handle the continuity of the solution across\n",
    "the interface $\\Gamma$ between $\\Omega_1$ and $\\Omega_2$.\n",
    "\n",
    "The resulting weak formulation is:\n",
    "$$\n",
    "\\text{find }u_1 \\in V(\\Omega_1), u_2 \\in V(\\Omega_2), \\eta \\in E(\\Gamma)\n",
    "$$\n",
    "s.t.\n",
    "$$\n",
    "\\int_{\\Omega_1} \\nabla u_1 \\cdot \\nabla v_1 dx +\n",
    "\\int_{\\Omega_2} \\nabla u_2 \\cdot \\nabla v_2 dx +\n",
    "\\int_{\\Gamma} \\lambda (v_1 - v_2) ds = 0,\n",
    "\\qquad \\forall v_1 \\in V(\\Omega_1), v_2 \\in V(\\Omega_2)\n",
    "$$\n",
    "and\n",
    "$$\n",
    "\\int_{\\Gamma} \\eta  (u_1 - u_2) ds = 0,\n",
    "\\qquad \\forall \\eta \\in E(\\Gamma)\n",
    "$$\n",
    "where boundary conditions on $\\partial\\Omega$ are embedded in $V(\\Omega_i) \\subset H^1(\\Omega_i)$, $i = 1, 2$, and $E(\\Gamma) \\subset L^2(\\Gamma)$.\n",
    "\n",
    "This example is a prototypical case of problems containing interface restricted variables (the Lagrange multiplier, in this case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import grad, inner, Measure, TestFunction, TrialFunction\n",
    "from dolfinx import DirichletBC, Function, FunctionSpace\n",
    "from dolfinx.cpp.mesh import GhostMode\n",
    "from dolfinx.fem import (LinearProblem, locate_dofs_topological)\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (assemble_matrix_block, assemble_scalar, assemble_vector_block,\n",
    "                               BlockVecSubVectorWrapper, create_vector_block, DofMapRestriction)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if MPI.COMM_WORLD.size > 1:\n",
    "    mesh_ghost_mode = GhostMode.shared_facet  # shared_facet ghost mode is required by dS\n",
    "else:\n",
    "    mesh_ghost_mode = GhostMode.none\n",
    "with XDMFFile(MPI.COMM_WORLD, \"data/circle.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh(mesh_ghost_mode)\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    subdomains = infile.read_meshtags(mesh, name=\"subdomains\")\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "cells_Omega1 = subdomains.indices[subdomains.values == 1]\n",
    "cells_Omega2 = subdomains.indices[subdomains.values == 2]\n",
    "facets_partial_Omega = boundaries.indices[boundaries.values == 1]\n",
    "facets_Gamma = boundaries.indices[boundaries.values == 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)\n",
    "dS = Measure(\"dS\")(subdomain_data=boundaries)\n",
    "dS = dS(2)  # restrict to the interface, which has facet ID equal to 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With domain decomposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function spaces\n",
    "V = FunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "V1 = V.clone()\n",
    "V2 = V.clone()\n",
    "M = V.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_V1_Omega1 = locate_dofs_topological(V1, subdomains.dim, cells_Omega1)\n",
    "dofs_V2_Omega2 = locate_dofs_topological(V2, subdomains.dim, cells_Omega2)\n",
    "dofs_M_Gamma = locate_dofs_topological(M, boundaries.dim, facets_Gamma)\n",
    "restriction_V1_Omega1 = DofMapRestriction(V1.dofmap, dofs_V1_Omega1)\n",
    "restriction_V2_Omega2 = DofMapRestriction(V2.dofmap, dofs_V2_Omega2)\n",
    "restriction_M_Gamma = DofMapRestriction(M.dofmap, dofs_M_Gamma)\n",
    "restriction = [restriction_V1_Omega1, restriction_V2_Omega2, restriction_M_Gamma]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "(u1, u2, l) = (TrialFunction(V1), TrialFunction(V2), TrialFunction(M))\n",
    "(v1, v2, m) = (TestFunction(V1), TestFunction(V2), TestFunction(M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "zero = Function(V)\n",
    "a = [[inner(grad(u1), grad(v1)) * dx(1), None, l(\"-\") * v1(\"-\") * dS],\n",
    "     [None, inner(grad(u2), grad(v2)) * dx(2), - l(\"+\") * v2(\"+\") * dS],\n",
    "     [m(\"-\") * u1(\"-\") * dS, - m(\"+\") * u2(\"+\") * dS, None]]\n",
    "f = [v1 * dx(1), v2 * dx(2), zero * m(\"-\") * dS]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define boundary conditions\n",
    "dofs_V1_partial_Omega = locate_dofs_topological((V1, V), boundaries.dim, facets_partial_Omega)\n",
    "dofs_V2_partial_Omega = locate_dofs_topological((V2, V), boundaries.dim, facets_partial_Omega)\n",
    "bc1 = DirichletBC(zero, dofs_V1_partial_Omega, V1)\n",
    "bc2 = DirichletBC(zero, dofs_V2_partial_Omega, V2)\n",
    "bcs = [bc1, bc2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system\n",
    "A = assemble_matrix_block(a, bcs=bcs, restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = assemble_vector_block(f, a, bcs=bcs, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u1u2l = create_vector_block(f, restriction=restriction)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, u1u2l)\n",
    "u1u2l.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "(u1, u2, l) = (Function(V1), Function(V2), Function(M))\n",
    "with BlockVecSubVectorWrapper(u1u2l, [V1.dofmap, V2.dofmap, M.dofmap], restriction) as u1u2l_wrapper:\n",
    "    for u1u2l_wrapper_local, component in zip(u1u2l_wrapper, (u1, u2, l)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = u1u2l_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u1, \"u1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u2, \"u2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, l, \"l\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without domain decomposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem forms\n",
    "a_ex = inner(grad(u), grad(v)) * dx\n",
    "f_ex = v * dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dirichlet BC object on Gamma\n",
    "dofs_V_partial_Omega = locate_dofs_topological(V, boundaries.dim, facets_partial_Omega)\n",
    "bc_ex = DirichletBC(zero, dofs_V_partial_Omega)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u_ex = Function(V)\n",
    "problem_ex = LinearProblem(\n",
    "    a_ex, f_ex, bcs=[bc_ex], u=u_ex,\n",
    "    petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\", \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "problem_ex.solve()\n",
    "u_ex.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_ex, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison and error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_ex1_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(u_ex, u_ex) * dx(1)), op=MPI.SUM))\n",
    "u_ex2_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(u_ex, u_ex) * dx(2)), op=MPI.SUM))\n",
    "err1_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(u_ex - u1, u_ex - u1) * dx(1)), op=MPI.SUM))\n",
    "err2_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(u_ex - u2, u_ex - u2) * dx(2)), op=MPI.SUM))\n",
    "print(\"Relative error on subdomain 1\", err1_norm / u_ex1_norm)\n",
    "print(\"Relative error on subdomain 2\", err2_norm / u_ex2_norm)\n",
    "assert np.isclose(err1_norm / u_ex1_norm, 0., atol=1.e-10)\n",
    "assert np.isclose(err2_norm / u_ex2_norm, 0., atol=1.e-10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
